"""
督导系统模型训练主脚本
整合风险预警、问题分类、热力图预测等模型的训练流程
"""

import logging
import argparse
import sys
import os
from datetime import datetime
import warnings
warnings.filterwarnings('ignore')

# 添加项目路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from data_processing.data_loader import DataLoader, DataPreprocessor
from models.risk_prediction_model import RiskModelTrainer
from models.problem_classification_model import ProblemClassificationModel
from models.heatmap_prediction_model import HeatmapPredictionModel
from config.model_config import config

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler(f'{config.data.log_dir}/training.log', encoding='utf-8'),
        logging.StreamHandler()
    ]
)

class SupervisionModelTrainer:
    """督导系统模型训练器"""
    
    def __init__(self):
        self.logger = logging.getLogger(__name__)
        self.data_loader = DataLoader()
        self.data_preprocessor = DataPreprocessor()
        
    def train_all_models(self, train_risk: bool = True, 
                        train_classification: bool = True,
                        train_heatmap: bool = True):
        """训练所有模型"""
        
        self.logger.info("开始训练督导系统所有模型...")
        start_time = datetime.now()
        
        # 加载数据
        self.logger.info("正在加载数据...")
        supervision_df = self.data_loader.load_supervision_records()
        area_df = self.data_loader.load_area_info()
        historical_df = self.data_loader.load_historical_problems()
        
        if supervision_df.empty:
            self.logger.error("未能加载督导记录数据，训练终止")
            return
        
        results = {}
        
        # 训练风险预警模型
        if train_risk:
            try:
                self.logger.info("="*50)
                self.logger.info("开始训练风险预警模型")
                risk_results = self._train_risk_prediction_model(
                    supervision_df, area_df, historical_df
                )
                results['risk_prediction'] = risk_results
                self.logger.info("风险预警模型训练完成")
            except Exception as e:
                self.logger.error(f"风险预警模型训练失败: {e}")
                results['risk_prediction'] = {'error': str(e)}
        
        # 训练问题分类模型
        if train_classification:
            try:
                self.logger.info("="*50)
                self.logger.info("开始训练问题分类模型")
                classification_results = self._train_problem_classification_model(supervision_df)
                results['problem_classification'] = classification_results
                self.logger.info("问题分类模型训练完成")
            except Exception as e:
                self.logger.error(f"问题分类模型训练失败: {e}")
                results['problem_classification'] = {'error': str(e)}
        
        # 训练热力图预测模型
        if train_heatmap:
            try:
                self.logger.info("="*50)
                self.logger.info("开始训练热力图预测模型")
                heatmap_results = self._train_heatmap_prediction_model(supervision_df)
                results['heatmap_prediction'] = heatmap_results
                self.logger.info("热力图预测模型训练完成")
            except Exception as e:
                self.logger.error(f"热力图预测模型训练失败: {e}")
                results['heatmap_prediction'] = {'error': str(e)}
        
        # 训练总结
        end_time = datetime.now()
        total_time = end_time - start_time
        
        self.logger.info("="*50)
        self.logger.info("所有模型训练完成")
        self.logger.info(f"总训练时间: {total_time}")
        
        # 生成训练报告
        self._generate_training_summary(results, total_time)
        
        return results
    
    def _train_risk_prediction_model(self, supervision_df, area_df, historical_df):
        """训练风险预警模型"""
        
        # 数据预处理
        features, targets = self.data_preprocessor.preprocess_for_risk_prediction(
            supervision_df, area_df, historical_df
        )
        
        if features.empty or targets.empty:
            raise ValueError("风险预警模型训练数据为空")
        
        # 创建训练器并训练
        trainer = RiskModelTrainer()
        results = trainer.train_and_evaluate(features, targets)
        
        # 记录关键指标
        evaluation = results['evaluation_results']
        self.logger.info(f"风险预警模型性能 - AUC: {evaluation['auc_score']:.4f}, "
                        f"准确率: {evaluation['accuracy']:.4f}, "
                        f"F1得分: {evaluation['f1_score']:.4f}")
        
        return results
    
    def _train_problem_classification_model(self, supervision_df):
        """训练问题分类模型"""
        
        # 数据预处理
        texts, problem_types, severity_levels = self.data_preprocessor.preprocess_for_problem_classification(
            supervision_df
        )
        
        if len(texts) < 10:
            raise ValueError("问题分类模型训练数据不足，至少需要10条记录")
        
        # 创建模型并训练
        model = ProblemClassificationModel()
        results = model.train(texts, problem_types, severity_levels)
        
        # 记录关键指标
        self.logger.info(f"问题分类模型性能 - "
                        f"问题类型准确率: {results['problem_type_accuracy']:.4f}, "
                        f"严重程度准确率: {results['severity_accuracy']:.4f}")
        
        return results
    
    def _train_heatmap_prediction_model(self, supervision_df):
        """训练热力图预测模型"""
        
        # 数据预处理
        X, y = self.data_preprocessor.preprocess_for_heatmap_prediction(supervision_df)
        
        if X.shape[0] < 50:
            raise ValueError("热力图预测模型训练数据不足，至少需要50个时间序列样本")
        
        # 创建模型并训练
        model = HeatmapPredictionModel()
        results = model.train(X, y)
        
        # 记录关键指标
        self.logger.info(f"热力图预测模型性能 - "
                        f"MSE: {results['mse']:.6f}, "
                        f"MAE: {results['mae']:.6f}, "
                        f"相关系数: {results['correlation']:.4f}")
        
        return results
    
    def _generate_training_summary(self, results, total_time):
        """生成训练总结报告"""
        
        summary = {
            'training_timestamp': datetime.now().isoformat(),
            'total_training_time': str(total_time),
            'model_results': results,
            'data_config': config.data.__dict__,
            'training_summary': {
                'models_trained': len([k for k, v in results.items() if 'error' not in v]),
                'models_failed': len([k for k, v in results.items() if 'error' in v]),
                'success_rate': len([k for k, v in results.items() if 'error' not in v]) / len(results) if results else 0
            }
        }
        
        # 保存总结报告
        import json
        summary_path = f"{config.data.log_dir}/training_summary_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
        with open(summary_path, 'w', encoding='utf-8') as f:
            json.dump(summary, f, ensure_ascii=False, indent=2, default=str)
        
        self.logger.info(f"训练总结报告已保存到: {summary_path}")
        
        # 打印总结信息
        self.logger.info("-" * 50)
        self.logger.info("训练总结:")
        for model_name, result in results.items():
            if 'error' in result:
                self.logger.info(f"❌ {model_name}: 训练失败 - {result['error']}")
            else:
                self.logger.info(f"✅ {model_name}: 训练成功")
        
        success_rate = summary['training_summary']['success_rate']
        self.logger.info(f"总体成功率: {success_rate:.1%}")

def create_sample_data():
    """创建示例数据用于测试"""
    
    import pandas as pd
    import numpy as np
    from datetime import datetime, timedelta
    
    # 创建示例督导记录数据
    np.random.seed(42)
    n_records = 1000
    
    area_codes = [f"A{str(i).zfill(3)}" for i in range(1, 21)]  # 20个区域
    problem_types = config.problem_classification.problem_types
    severity_levels = config.problem_classification.severity_levels
    
    # 生成时间序列
    start_date = datetime.now() - timedelta(days=365)
    dates = [start_date + timedelta(days=i) for i in range(365)]
    
    records = []
    for i in range(n_records):
        records.append({
            'id': i + 1,
            'supervision_date': np.random.choice(dates),
            'area_code': np.random.choice(area_codes),
            'area_name': f"区域{np.random.choice(area_codes)}",
            'longitude': 116.4 + np.random.normal(0, 0.1),
            'latitude': 39.9 + np.random.normal(0, 0.1),
            'problem_description': f"发现{np.random.choice(problem_types)}问题，需要{np.random.choice(['立即', '尽快', '按计划'])}处理",
            'problem_type': np.random.choice(problem_types),
            'severity_level': np.random.choice(severity_levels),
            'supervisor_id': f"SUP{np.random.randint(1, 51):03d}",
            'rectification_deadline': np.random.choice(dates),
            'rectification_status': np.random.choice(['进行中', '已完成', '延期']),
            'follow_up_date': np.random.choice(dates),
            'created_at': datetime.now(),
            'updated_at': datetime.now()
        })
    
    supervision_df = pd.DataFrame(records)
    
    # 创建区域信息数据
    area_records = []
    for area_code in area_codes:
        area_records.append({
            'area_code': area_code,
            'area_name': f"区域{area_code}",
            'area_type': np.random.choice(['工业区', '住宅区', '商业区', '混合区']),
            'parent_area_code': None,
            'longitude': 116.4 + np.random.normal(0, 0.1),
            'latitude': 39.9 + np.random.normal(0, 0.1),
            'population': np.random.randint(1000, 50000),
            'area_size': np.random.uniform(1.0, 20.0),
            'risk_level': np.random.choice(['低', '中', '高']),
            'management_unit': f"管理单位{np.random.randint(1, 11)}"
        })
    
    area_df = pd.DataFrame(area_records)
    
    return supervision_df, area_df

def main():
    """主函数"""
    
    parser = argparse.ArgumentParser(description='训练督导系统模型')
    parser.add_argument('--models', nargs='+', 
                       choices=['risk', 'classification', 'heatmap', 'all'],
                       default=['all'], help='要训练的模型类型')
    parser.add_argument('--data-dir', type=str, default='data',
                       help='数据目录，可以使用 real_data 目录')
    parser.add_argument('--use-sample-data', action='store_true',
                       help='使用示例数据进行训练')
    parser.add_argument('--config-file', type=str, default=None,
                       help='自定义配置文件路径')
    
    args = parser.parse_args()
    
    # 加载自定义配置
    if args.config_file:
        # 这里可以加载自定义配置文件
        pass
    
    # 确定要训练的模型
    train_risk = 'all' in args.models or 'risk' in args.models
    train_classification = 'all' in args.models or 'classification' in args.models
    train_heatmap = 'all' in args.models or 'heatmap' in args.models
    
    # 创建训练器
    trainer = SupervisionModelTrainer()
    
    # 更新配置中的数据目录
    if args.data_dir != 'data':
        config.data.data_dir = args.data_dir
        trainer.logger.info(f"使用数据目录: {args.data_dir}")
    
    # 如果使用示例数据，创建示例数据文件
    if args.use_sample_data:
        trainer.logger.info("创建示例数据用于训练...")
        supervision_df, area_df = create_sample_data()
        
        # 保存示例数据
        os.makedirs(config.data.data_dir, exist_ok=True)
        supervision_df.to_csv(f'{config.data.data_dir}/sample_supervision_records.csv', index=False)
        area_df.to_csv(f'{config.data.data_dir}/sample_area_info.csv', index=False)
        trainer.logger.info("示例数据已创建并保存")
    
    try:
        # 开始训练
        results = trainer.train_all_models(
            train_risk=train_risk,
            train_classification=train_classification,
            train_heatmap=train_heatmap
        )
        
        print("\n" + "="*60)
        print("🎉 模型训练完成！")
        print("="*60)
        
        # 显示结果摘要
        for model_name, result in results.items():
            if 'error' in result:
                print(f"❌ {model_name}: 训练失败")
            else:
                print(f"✅ {model_name}: 训练成功")
        
        print(f"\n📊 模型和日志文件保存在:")
        print(f"   模型目录: {config.data.model_dir}")
        print(f"   日志目录: {config.data.log_dir}")
        
    except Exception as e:
        trainer.logger.error(f"训练过程中发生错误: {e}")
        print(f"\n❌ 训练失败: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main() 